# @Time : 2021/8/5 11:22
# @Author : Li Kunlun
# @Description : 卷积神经网络（LeNet）


import utils as d2l
import mxnet as mx
from mxnet import autograd, gluon, init, nd
from mxnet.gluon import loss as gloss, nn
import time

# 1、LeNet模型
net = nn.Sequential()
net.add(nn.Conv2D(channels=6, kernel_size=5, activation='sigmoid'),
        nn.MaxPool2D(pool_size=2, strides=2),
        nn.Conv2D(channels=16, kernel_size=5, activation='sigmoid'),
        nn.MaxPool2D(pool_size=2, strides=2),
        # Dense会默认将(批量大小, 通道, 高, 宽)形状的输入转换成(批量大小, 通道 * 高 * 宽)形状的输入
        nn.Dense(120, activation='sigmoid'),
        nn.Dense(84, activation='sigmoid'),
        nn.Dense(10))

# 构造一个高和宽均为28的单通道数据样本，并逐层进行前向计算来查看每个层的输出形状
X = nd.random.uniform(shape=(1, 1, 28, 28))
net.initialize()
for layer in net:
    X = layer(X)
    """
    conv0 output shape:	 (1, 6, 24, 24)
    pool0 output shape:	 (1, 6, 12, 12)
    conv1 output shape:	 (1, 16, 8, 8)
    pool1 output shape:	 (1, 16, 4, 4)
    dense0 output shape:	 (1, 120)
    dense1 output shape:	 (1, 84)
    dense2 output shape:	 (1, 10)
    """
    print(layer.name, 'output shape:\t', X.shape)

# 2、获取数据和训练模型
# 使用Fashion-MNIST作为训练数据集
# batch_size = 256
batch_size = 10
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size=batch_size)


# 因为卷积神经网络计算比多层感知机要复杂，建议使用GPU来加速计算。
def try_gpu():
    try:
        print("-----gpu1-----------")
        ctx = mx.cpu()
        # ctx = mx.gpu()
        # _ = nd.zeros((1,), ctx=ctx)
        print("-----gpu2-----------")
    except mx.base.MXNetError:
        ctx = mx.cpu()
        print("-----cpu-----------")
    return ctx


ctx = try_gpu()
# gpu(0)
print(ctx)


# 计算准确率
def evaluate_accuracy(data_iter, net, ctx):
    acc_sum, n = nd.array([0], ctx=ctx), 0
    for X, y in data_iter:
        # 如果ctx代表GPU及相应的显存，将数据复制到显存上
        X, y = X.as_in_context(ctx), y.as_in_context(ctx).astype('float32')
        acc_sum += (net(X).argmax(axis=1) == y).sum()
        n += y.size
    return acc_sum.asscalar() / n


# 训练模型
def train_ch5(net, train_iter, test_iter, batch_size, trainer, ctx, num_epochs):
    print('training on', ctx)
    loss = gloss.SoftmaxCrossEntropyLoss()
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n, start = 0.0, 0.0, 0, time.time()
        for X, y in train_iter:
            X, y = X.as_in_context(ctx), y.as_in_context(ctx)
            with autograd.record():
                y_hat = net(X)
                l = loss(y_hat, y).sum()
            l.backward()
            trainer.step(batch_size)
            y = y.astype('float32')
            train_l_sum += l.asscalar()
            train_acc_sum += (y_hat.argmax(axis=1) == y).sum().asscalar()
            n += y.size
        test_acc = evaluate_accuracy(test_iter, net, ctx)
        print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f, ''time %.1f sec' % (
        epoch + 1, train_l_sum / n, train_acc_sum / n, test_acc, time.time() - start))


# 测试
lr, num_epochs = 0.9, 5
net.initialize(force_reinit=True, ctx=ctx, init=init.Xavier())
trainer = gluon.Trainer(net.collect_params(), 'sgd', {'learning_rate': lr})
train_ch5(net, train_iter, test_iter, batch_size, trainer, ctx, num_epochs)
